import pandas as pd
from sentence_transformers import SentenceTransformer
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import os
os.environ['TRANSFORMERS_CACHE'] = './models'

# os.environ['HTTP_PROXY'] = 'http://127.0.0.1:21882'
# os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:21882'
import logging
logging.basicConfig(level=logging.INFO)

class ExcelKnowledgeBase:
    def __init__(self, excel_path):
        self.df = pd.read_excel(excel_path)
        print('-----------模型开始初始化-----------')
        self.encoder = SentenceTransformer('./models/models--sentence-transformers--paraphrase-multilingual-MiniLM-L12-v2/snapshots/86741b4e3f5cb7765a600d3a3d55a0f6a6cb443d',device='cpu')
        print('-----------模型初始化结束-----------')
        # 预编码所有问题
        self.questions = self.df['问题描述'].tolist()
        self.question_embeddings = self.encoder.encode(self.questions)

    def find_most_similar(self, query, threshold=0.6):
        # 编码用户查询
        query_embedding = self.encoder.encode([query])

        # 计算相似度
        similarities = cosine_similarity(query_embedding, self.question_embeddings)[0]
        max_index = np.argmax(similarities)
        max_similarity = similarities[max_index]
        print(max_similarity,'++++++')
        if max_similarity > threshold:
            return self.df.iloc[max_index]['处理方案/问题答复']
        return None

if __name__ == '__main__':
    print("--------------hello----------------------")
    base = ExcelKnowledgeBase("D:\\项目\\智能客服\\慕卓柯BMS知识库整理.xlsx")
    print(base.find_most_similar("登录地址是什么？"))
    print(base.find_most_similar("修改密码？"))